package cn.iocoder.yudao.module.security.service.dashboard;

import cn.hutool.core.date.DatePattern;
import cn.hutool.core.date.LocalDateTimeUtil;
import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
import cn.iocoder.yudao.module.security.controller.admin.dashboard.vo.*;
import cn.iocoder.yudao.module.security.dal.dataobject.asset.SecurityAssetDO;
import cn.iocoder.yudao.module.security.dal.dataobject.whitelistevent.WhitelistEventEsDO;
import cn.iocoder.yudao.module.security.dal.mysql.asset.SecurityAssetMapper;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.FieldValue;
import co.elastic.clients.elasticsearch._types.SortOrder;
import co.elastic.clients.elasticsearch._types.Time;
import co.elastic.clients.elasticsearch._types.aggregations.*;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.elasticsearch.core.search.HitsMetadata;
import co.elastic.clients.util.NamedValue;
import com.alibaba.fastjson.JSON;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.IOException;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.stream.Collectors;

@Slf4j
@Service
public class DashboardService {

    @Resource
    private ElasticsearchClient elasticsearchClient;

    @Resource
    private SecurityAssetMapper securityAssetMapper;

    @Resource
    private SecurityCommonService securityCommonService;

    public List<SecurityDashboardNetworkRespVO> getNetwork(SecurityDashboardReqVO reqVO) {

        List<SecurityDashboardNetworkRespVO> list = new ArrayList<>();
        // 构建复合聚合
        List<Map<String, CompositeAggregationSource>> sources = new ArrayList<>();
        Map<String, CompositeAggregationSource> stk1 = new HashMap<>();
        stk1.put("stk1", new CompositeAggregationSource.Builder().terms(t -> t.field("destination.ip")).build());
        Map<String, CompositeAggregationSource> stk2 = new HashMap<>();
        stk2.put("stk2", new CompositeAggregationSource.Builder().terms(t -> t.field("source.ip")).build());
        sources.add(stk1);
        sources.add(stk2);
        CompositeAggregation compositeAggregation = AggregationBuilders.composite()
                .sources(sources).size(10000)
                .build();

        Result result = getResult(reqVO);

        // 构建搜索请求
        SearchRequest request = SearchRequest.of(s -> s
                .index("*-pfelk-snort*") // 原查询的索引模式
                .query(q -> q
                        .bool(b -> {
                            b.filter(f -> f.range(r -> r.date(a -> a.field("@timestamp")
                                    .gte(result.getStartTimeStr())
                                    .lte(result.getEndTimeStr()))));
                            return b;
                        })
                )
                .size(0) // 不返回具体文档
                .aggregations("table", compositeAggregation._toAggregation()) // 添加复合聚合
        );

        try {
            // 执行查询
            SearchResponse<Object> response = elasticsearchClient.search(request, Object.class);

            // 处理聚合结果（示例处理逻辑）
            CompositeAggregate tableAgg = response.aggregations()
                    .get("table")
                    .composite();

            List<SecurityAssetDO> assetDOS = securityAssetMapper.selectList(
                    new LambdaQueryWrapperX<SecurityAssetDO>()
                            .eq(SecurityAssetDO::getStatus, CommonStatusEnum.ENABLE.getStatus())
                            .eq(SecurityAssetDO::getDeleted, false)
            );
            Map<String, String> assetMap = assetDOS.stream()
                    .filter(assetDO -> StringUtils.isNotBlank(assetDO.getIpAddress()) && StringUtils.isNotBlank(assetDO.getRemark()))
                    .collect(Collectors.toMap(SecurityAssetDO::getIpAddress, SecurityAssetDO::getRemark));
            for (CompositeBucket bucket : tableAgg.buckets().array()) {
                String destIp = bucket.key().get("stk1").stringValue();   // 目的IP
                String srcIp = bucket.key().get("stk2").stringValue();   // 源IP
                long docCount = bucket.docCount();                     // 计数
                SecurityDashboardNetworkRespVO vo = new SecurityDashboardNetworkRespVO();
                vo.setSourceIp(srcIp);
                vo.setDescIp(destIp);
                vo.setCount(docCount);
                vo.setAssetName(assetMap.get(srcIp));
                list.add(vo);
            }


        } catch (Exception e) {
            log.error("elasticsearch request err,request: {}", request, e);
        }
        return list;
    }

    private static Result getResult(SecurityDashboardReqVO reqVO) {
        long now = System.currentTimeMillis();
        long dayAgo14 = now - 14 * 24 * 60 * 60 * 1000;
        LocalDateTime startTime = LocalDateTimeUtil.of(dayAgo14);
        LocalDateTime endTime = LocalDateTimeUtil.of(now);
        if (reqVO.getCreateTime() != null && reqVO.getCreateTime().length == 2) {
            startTime = LocalDateTimeUtil.of(reqVO.getCreateTime()[0]);
            endTime = LocalDateTimeUtil.of(reqVO.getCreateTime()[1]);
        }

        String startTimeStr = LocalDateTimeUtil.format(startTime, DatePattern.UTC_SIMPLE_PATTERN) + "+08:00";
        String endTimeStr = LocalDateTimeUtil.format(endTime, DatePattern.UTC_SIMPLE_PATTERN) + "+08:00";
        Result result = new Result(startTimeStr, endTimeStr);
        return result;
    }

    @Data
    public static class Result {
        public String startTimeStr;
        public String endTimeStr;

        public Result(String startTimeStr, String endTimeStr) {
            this.startTimeStr = startTimeStr;
            this.endTimeStr = endTimeStr;
        }
    }

    public List<SecurityDashboardProtocolRespVO> getProtocolApp(SecurityDashboardReqVO reqVO) {
        // 4. 构建聚合分析
        TermsAggregation termsAggregation = AggregationBuilders.terms()
                .field("network.app")
                .size(20)
                .order(new NamedValue<SortOrder>("_count", SortOrder.Desc))
                .build();

        Result result = getResult(reqVO);

        SearchRequest request = SearchRequest.of(s -> s
                .index("*-pfelk-snort*")
                .query(q -> q
                        .bool(b -> {
                            b.filter(f -> f.range(r -> r.date(a -> a.field("@timestamp")
                                    .gte(result.getStartTimeStr())
                                    .lte(result.getEndTimeStr()))));
                            return b;
                        })
                )
                .aggregations("network_app", termsAggregation._toAggregation())
                .size(0)
        );


        List<SecurityDashboardProtocolRespVO> list = new ArrayList<>();
        // 执行并处理结果
        try {
            SearchResponse<Object> response = elasticsearchClient.search(request, Object.class);

            // 7. 解析聚合结果
            List<StringTermsBucket> buckets = response.aggregations()
                    .get("network_app")
                    .sterms()
                    .buckets()
                    .array();
            for (StringTermsBucket bucket : buckets) {
                String key = bucket.key().stringValue();
                long count = bucket.docCount();
                SecurityDashboardProtocolRespVO vo = new SecurityDashboardProtocolRespVO();
                vo.setProtocol(key);
                vo.setCount(count);
                list.add(vo);
            }
        } catch (IOException e) {
            log.error("elasticsearch request err,request: {}", request, e);
        }

        return list;
    }

    public List<SecurityDashboardProtocolRespVO> getProtocolType(SecurityDashboardReqVO reqVO) {
        // 4. 构建聚合分析
        TermsAggregation termsAggregation = AggregationBuilders.terms()
                .field("network.protocol")
                .size(20)
                .order(new NamedValue<SortOrder>("_count", SortOrder.Desc))
                .build();
        SearchRequest request = SearchRequest.of(s -> s
                .index("*-pfelk-firewall*")
                .aggregations("network_protocol", termsAggregation._toAggregation())
                .size(0)
        );
        List<SecurityDashboardProtocolRespVO> list = new ArrayList<>();
        // 执行并处理结果
        try {
            SearchResponse<Object> response = elasticsearchClient.search(request, Object.class);

            // 7. 解析聚合结果
            List<StringTermsBucket> buckets = response.aggregations()
                    .get("network_protocol")
                    .sterms()
                    .buckets()
                    .array();
            Map<String, Long> map = new HashMap<>();
            for (StringTermsBucket bucket : buckets) {
                String key = bucket.key().stringValue();
                long count = bucket.docCount();
                key = key.toUpperCase();
                if (map.containsKey(key)) {
                    count += map.get(key);
                }
                map.put(key, count);
            }
            for (Map.Entry<String, Long> entry : map.entrySet()) {
                String key = entry.getKey();
                long count = entry.getValue();
                SecurityDashboardProtocolRespVO vo = new SecurityDashboardProtocolRespVO();
                vo.setProtocol(key);
                vo.setCount(count);
                list.add(vo);
            }
        } catch (IOException e) {
            log.error("elasticsearch request err,request: {}", request, e);
        }

        return list;
    }

    public List<SecurityDashboardProtocolBarRespVO> getProtocolBar(SecurityDashboardReqVO reqVO) {

        long now = System.currentTimeMillis();
        long dayAgo21 = now - 20 * 24 * 60 * 60 * 1000;

        String startTimeStr = LocalDateTimeUtil.format(LocalDateTimeUtil.of(dayAgo21), DatePattern.UTC_SIMPLE_PATTERN) + "+08:00";
        String endTimeStr = LocalDateTimeUtil.format(LocalDateTimeUtil.of(now), DatePattern.UTC_SIMPLE_PATTERN) + "+08:00";


        // 构建日期直方图子聚合
        DateHistogramAggregation dateHistogramAgg = AggregationBuilders.dateHistogram()
                .field("@timestamp")
                //.fixedInterval("12h")
                .fixedInterval(new Time.Builder().time("12h").build())
                .timeZone("Asia/Shanghai")
                /*.extendedBounds(b -> b
                        .min(1740117141346L)
                        .max(1741931541346L)
                )*/
                .extendedBounds(
                        new ExtendedBounds.Builder<FieldDateMath>()
                                .min(new FieldDateMath.Builder().value((double) dayAgo21).build())
                                .max(new FieldDateMath.Builder().value((double) now).build())
                                .build()
                )
                .build();
        /*new TermsAggregationBuilder("network_app").field("network.app")
                .size(5)
                .shardSize(25)
                //.order(List.of(new NamedSortOrder("_count", SortOrder.Desc)))
                .order(new InternalOrder.("network_app", SortOrder.Desc) {
                })*/
        //area_field_builder.su
        // 构建主terms聚合
        /*TermsAggregation termsAggregation = AggregationBuilders.terms()
                .field("network.app")
                .size(5)
                .shardSize(25)
                //.order(List.of(new NamedSortOrder("_count", SortOrder.Desc)))
                .order(new NamedValue<SortOrder>("_count", SortOrder.Desc))
                //.aggregations("1", dateHistogramAgg._toAggregation()) // 添加子聚合
                .build();*/
        Aggregation termsAggregation = new Aggregation.Builder()
                .terms(t -> t
                        .field("network.app")
                        .size(5)
                        .shardSize(25)
                        //.order(List.of(new NamedSortOrder("_count", SortOrder.Desc)))
                        .order(new NamedValue<SortOrder>("_count", SortOrder.Desc))
                )
                .aggregations("1", dateHistogramAgg._toAggregation()) // 关键点：这里添加子聚合
                .build();

        // 构建搜索请求
        SearchRequest request = SearchRequest.of(s -> s
                .index("*-pfelk-snort*")
                .size(0)
                .query(q -> q
                        .bool(b -> b
                                .filter(
                                        f -> f.range(r -> r.date(a -> a.field("@timestamp")
                                                .gte(startTimeStr)
                                                .lte(endTimeStr)))
                                )
                        )
                )
                .aggregations("0", termsAggregation)
        );

        List<SecurityDashboardProtocolBarRespVO> resultList = new ArrayList<>();

        try {
            SearchResponse<Object> response = elasticsearchClient.search(request, Object.class);

            // 处理主聚合结果
            StringTermsAggregate mainAgg = response.aggregations()
                    .get("0")
                    .sterms();

            for (StringTermsBucket bucket : mainAgg.buckets().array()) {
                String protocol = bucket.key().stringValue();
                long totalCount = bucket.docCount();

                // 处理子聚合日期直方图
                DateHistogramAggregate dateHistogram = bucket.aggregations().get("1").dateHistogram();
                List<DateHistogramBucket> timeBuckets = dateHistogram.buckets().array();

                // 构建响应对象（示例）
                SecurityDashboardProtocolBarRespVO vo = new SecurityDashboardProtocolBarRespVO();
                vo.setProtocol(protocol);
                vo.setCount(totalCount);
                vo.setTimeDistribution(timeBuckets.stream()
                        .collect(Collectors.toMap(
                                DateHistogramBucket::key,
                                DateHistogramBucket::docCount
                        )));
                resultList.add(vo);
            }
        } catch (IOException e) {
            log.error("Elasticsearch query failed: {}", request, e);
        }

        return resultList;
    }

    public List<NetworkEventRespVO> getNetworkEvent(SecurityDashboardReqVO reqVO) {
        // 构建复合查询
        List<FieldValue> list = securityCommonService.buildProtocolType();
        List<FieldValue> eventTypeList = securityCommonService.buildEventType();

        SearchRequest request = SearchRequest.of(s -> s
                        .index("*-pfelk-snort*")
                        .query(q -> q
                                .bool(b -> {
                                    b.filter(f -> f.terms(t -> t.field("network.app")
                                            .terms(v -> v.value(list))));
                                            /*.filter(f -> f.range(r -> r.date(a -> a.field("@timestamp")
                                                    .gte(startTimeStr)
                                                    .lte(endTimeStr))));*/
                                    //.filter(f -> f.term(t -> t.field("elastalert.alert").value(0)))
                                    b.filter(f -> f.terms(t -> t.field("vulnerability.classification")
                                            .terms(v -> v.value(eventTypeList))));
                                    //b.filter(f -> f.term(t -> t.field("vulnerability.classification").value("关键事件告警")));
                                    return b;
                                })
                        )
                        .size(30)
                        .sort(so -> so.field(f -> f.field("@timestamp").order(SortOrder.Desc)))
                //.sort(so -> so.field(f -> f.field("event.sequence").order(SortOrder.Asc)))
        );

        // 执行并处理结果
        try {
            SearchResponse<Map> response = elasticsearchClient.search(request, Map.class);
            HitsMetadata<Map> hits = response.hits();
            long total = hits.total().value();
            List<NetworkEventRespVO> list1 = new ArrayList<>();
            for (Hit<Map> hit : hits.hits()) {
                String sourceStr = JSON.toJSONString(hit.source());
                //WhitelistEventEsDO source = JSON.parseObject(sourceStr, WhitelistEventEsDO.class);
                NetworkEventRespVO.NetworkEventRespVOBuilder builder = NetworkEventRespVO.builder();
                WhitelistEventEsDO.SourceData sourceData = JSON.parseObject(sourceStr, WhitelistEventEsDO.SourceData.class);
                if (Objects.nonNull(sourceData)) {
                    if (Objects.nonNull(sourceData.getSource())) {
                        builder.sourcePort(Integer.valueOf(sourceData.getSource().getPort()));
                        builder.sourceIp(sourceData.getSource().getIp());
                    }
                    if (Objects.nonNull(sourceData.getDestination())) {
                        builder.destinationPort(Integer.valueOf(sourceData.getDestination().getPort()));
                        builder.destinationIp(sourceData.getDestination().getIp());
                    }
                    if (Objects.nonNull(sourceData.getVulnerability())) {
                        builder.eventType(sourceData.getVulnerability().getClassification());
                    }
                    if (Objects.nonNull(sourceData.getNetwork())) {
                        builder.protocol(sourceData.getNetwork().getApp());
                    }
                    String origin = sourceData.getTimestamp();
                    int dotIndex = origin.indexOf('.');
                    String result = origin.substring(0, dotIndex);
                    LocalDateTime timestamp = LocalDateTime.parse(result, DateTimeFormatter.ofPattern("yyyy-MM-dd'T'HH:mm:ss"));
                    // 加上8个小时
                    timestamp = timestamp.plusHours(8);
                    builder.eventTime(timestamp);
                }

                NetworkEventRespVO build = builder.build();
                list1.add(build);
            }
            return list1;
        } catch (IOException e) {
            log.error("elasticsearch request err,request: {}", request, e);
        }

        return Collections.emptyList();
    }

    public List<SecurityDashboardProtocolRespVO> getNetworkEventType(SecurityDashboardReqVO reqVO) {
        // 4. 构建聚合分析
        TermsAggregation termsAggregation = AggregationBuilders.terms()
                .field("vulnerability.classification")
                .size(20)
                .order(new NamedValue<SortOrder>("_count", SortOrder.Desc))
                .build();

        List<FieldValue> eventTypeList = securityCommonService.buildEventType();

        Result result = getResult(reqVO);

        SearchRequest request = SearchRequest.of(s -> s
                .index("*-pfelk-snort*")
                .query(q -> q
                        .bool(b -> {
                            b.filter(f -> f.range(r -> r.date(a -> a.field("@timestamp")
                                    .gte(result.getStartTimeStr())
                                    .lte(result.getEndTimeStr()))));
                            //.filter(f -> f.term(t -> t.field("elastalert.alert").value(0)))
                            b.filter(f -> f.terms(t -> t.field("vulnerability.classification")
                                    .terms(v -> v.value(eventTypeList))));
                            return b;
                        })
                )
                .aggregations("network_event", termsAggregation._toAggregation())
                .size(0)
        );
        List<SecurityDashboardProtocolRespVO> list = new ArrayList<>();
        // 执行并处理结果
        try {
            String str = request.toString();
            log.info("elasticsearch request: {}", str);
            SearchResponse<Object> response = elasticsearchClient.search(request, Object.class);

            // 7. 解析聚合结果
            List<StringTermsBucket> buckets = response.aggregations()
                    .get("network_event")
                    .sterms()
                    .buckets()
                    .array();
            for (StringTermsBucket bucket : buckets) {
                String key = bucket.key().stringValue();
                long count = bucket.docCount();
                SecurityDashboardProtocolRespVO vo = new SecurityDashboardProtocolRespVO();
                vo.setProtocol(key);
                vo.setCount(count);
                list.add(vo);
            }
            Set<String> collect = eventTypeList.stream().map(FieldValue::stringValue).collect(Collectors.toSet());
            Set<String> collect1 = list.stream().map(SecurityDashboardProtocolRespVO::getProtocol).collect(Collectors.toSet());
            for (String s : collect) {
                if (!collect1.contains(s)) {
                    SecurityDashboardProtocolRespVO vo = new SecurityDashboardProtocolRespVO();
                    vo.setProtocol(s);
                    vo.setCount(0L);
                    list.add(vo);
                }
            }
        } catch (IOException e) {
            log.error("elasticsearch request err,request: {}", request, e);
        }

        return list;
    }
}
